from llama_index.core import Settings
from llama_index.core.agent import ReActAgent
from indexing import query_tool
import streamlit as st
import pandas as pd
import json, re
from dotenv import load_dotenv
from llama_index.core.chat_engine import SimpleChatEngine

from prompts import llm_prompt_testcase, agent_prompt_testcase, llm_prompt_qa, agent_prompt_qa

load_dotenv()

def validate_cases(cases):
    required_fields = ["用例编号", "步骤描述", "预期结果", "优先级"]
    for case in cases:
        if not all(field in case for field in required_fields):
            return False
    return True

def extract_json(response_str):
    # 使用正则匹配json代码块
    json_match = re.search(r'```json\n(.*?)\n```', response_str, re.DOTALL)
    if json_match:
        print(json_match)
        return json_match.group(1)
    return response_str


def generate_ai_response(user_input, model_name, enable_knowledge=True):
    # 动态设置模型
    ai_result = []
    from model_config import dashscope_llm
    Settings.llm = dashscope_llm(model_name)

    print(f"==========================当前模型：{model_name}============================")
    sources = []
    query_engine_tools = []
    try:
        tool1 = query_tool(index_id="project_docs",
                           description="包含当前项目的需求文档、技术方案、历史缺陷报告、项目特有约束条件等上下文信息。当需要结合项目特殊要求或规避已知风险时优先使用。")
        if tool1:
            query_engine_tools.append(tool1)
        tool2 = query_tool(index_id="testing_standards",
                           description="提供标准测试方法论、常见测试场景模版、行业测试规范、通用用例集等测试领域知识。当需要参考测试最佳实践或验证测试设计合理性时使用。")
        if tool2:
            query_engine_tools.append(tool2)
        tool3 = query_tool(index_id="business_rules",
                           description="存储业务流程规则、领域术语字典、用户旅程地图、业务关键性等级定义等业务知识。当涉及复杂业务逻辑验证或业务合规性检查时必须调用。")
        if tool3:
            query_engine_tools.append(tool3)

        if len(query_engine_tools) == 0 or not enable_knowledge:
            print("启用普通聊天模式")
            engine = SimpleChatEngine.from_defaults()
            prompt = llm_prompt_qa + f"需求描述如下：\n{user_input}"
            response = engine.chat(prompt)
            sources = []

        else:
            print("启用ReAct Agent")
            # 创建一个Agent
            agent = ReActAgent.from_tools(query_engine_tools, verbose=True)
            prompt = agent_prompt_qa + f"需求描述如下：\n{user_input}"
            response = agent.chat(prompt)
            # 输出查询到的文本块及其来源文件
            sources = []
            for node in response.source_nodes:
                src = {"文件名称": node.node.metadata['file_name'], "内容": node.node.get_content()}
                sources.append(src)

        ai_result = response
        ai_result =  json.loads(extract_json(str(response)))

    except json.JSONDecodeError:
        return ai_result,sources
    except Exception as e:
        st.error(f"发生错误：{e}")
        return [], []
    return ai_result,sources


def generate_test_cases(user_input, model_name="qwen-max", enable_knowledge=True):
    """生成测试用例"""
    # 动态设置模型
    from model_config import dashscope_llm
    Settings.llm = dashscope_llm(model_name)

    print(f"==========================当前模型：{model_name}============================")

    query_engine_tools = []
    try:
        tool1 = query_tool(index_id="project_docs",
                           description="包含当前项目的需求文档、技术方案、历史缺陷报告、项目特有约束条件等上下文信息。当需要结合项目特殊要求或规避已知风险时优先使用。")
        if tool1:
            query_engine_tools.append(tool1)
        tool2 = query_tool(index_id="testing_standards",
                           description="提供标准测试方法论、常见测试场景模版、行业测试规范、通用用例集等测试领域知识。当需要参考测试最佳实践或验证测试设计合理性时使用。")
        if tool2:
            query_engine_tools.append(tool2)
        tool3 = query_tool(index_id="business_rules",
                           description="存储业务流程规则、领域术语字典、用户旅程地图、业务关键性等级定义等业务知识。当涉及复杂业务逻辑验证或业务合规性检查时必须调用。")
        if tool3:
            query_engine_tools.append(tool3)

        if len(query_engine_tools) == 0 or not enable_knowledge:
            print("启用普通聊天模式")
            engine = SimpleChatEngine.from_defaults()
            prompt = llm_prompt_testcase+f"需求描述如下：\n{user_input}"
            response = engine.chat(prompt)
            sources = []
        else:
            print("启用ReAct Agent")
            # 创建一个Agent
            agent = ReActAgent.from_tools(query_engine_tools, verbose=True)
            prompt = agent_prompt_testcase+f"需求描述如下：\n{user_input}"
            response = agent.chat(prompt)
            # 输出查询到的文本块及其来源文件
            sources = []
            for node in response.source_nodes:
                src = { "文件名称": node.node.metadata['file_name'], "内容": node.node.get_content()}
                sources.append(src)

        cases = json.loads(extract_json(str(response)))
        if not validate_cases(cases):
            st.error("生成的测试用例格式不符合要求，无法解析")
            return [], []
    except json.JSONDecodeError:
        st.error("大模型的响应不是有效的JSON格式，无法被解析")
        return [], []
    except Exception as e:
        st.error(f"发生错误：{e}")
        return [], []
    return cases, sources


def display_results(cases, sources):
    """结果展示组件"""
    if not cases:
        return

    tab1, tab2, tab3 = st.tabs(["表格视图", "JSON数据", "参考片段"])

    with tab1:
        df = pd.DataFrame(cases)
        st.dataframe(df, use_container_width=True)

    with tab2:
        st.code(json.dumps(cases, indent=2, ensure_ascii=False))

    with tab3:
        if len(sources) == 0:
            st.html("<h1>没有参考文本<h1>")
        else:
            for src in sources:
                st.markdown(f"### {src['文件名称']}")
                st.code(src['内容'], language='``')
                st.markdown("")

def display_res(ai_result,sources):
    """普通结果展示组件"""

    tab1, tab2 = st.tabs(["AI响应结果",  "参考片段"])
    with tab1:
        # df = pd.DataFrame(sources)
        # st.dataframe(df, use_container_width=True)
        st.markdown(ai_result)
    with tab2:
        if len(sources) == 0:
            st.html("<h1>没有参考文本<h1>")
        else:
            for src in sources:
                st.markdown(f"### {src['文件名称']}")
                st.code(src['内容'], language='``')
                st.markdown("")


def main():
    st.set_page_config(page_title="AI测试助手", page_icon="📄", layout="wide")
    page = st.sidebar.radio("选择页面", ["测试用例生成", "知识库管理","行情答疑助手"])
    # 添加模型切换
    from model_config import DEFINED_MODELS
    selected_model = st.sidebar.selectbox(
        "选择大模型",
        options=list(DEFINED_MODELS.keys()),
        index=0  # 默认选中第一个
    )


    if page == "知识库管理":
        import knowledge_base
        knowledge_base.main()
    elif page == "行情答疑助手":
        st.title("📄 行情答疑助手")
        with st.form(key="main_form"):
            user_input = st.text_area("问题描述", height=150,
                                      placeholder="在此处描述您的问题")
            # 在main()函数的with st.form块中修改：
            col1, col2 = st.columns([3, 1])
            with col1:
                submitted = st.form_submit_button("一键问答")
            with col2:
                enable_knowledge = st.checkbox("是否参考知识库内容", value=False,
                                               help="启用后将会结合知识库内容生成更精准的回答")

        # 结果生成
        if submitted and user_input:
            with st.spinner("大模型正在努力思考中，请耐心等待哦..."):
                ai_result,sources = generate_ai_response(user_input, model_name=selected_model,
                                               enable_knowledge=enable_knowledge)
                print(ai_result,sources)
                display_res(ai_result,sources)
    else:
        st.title("📄 AI测试助手")
        with st.form(key="main_form"):
            user_input = st.text_area("需求描述", height=150,
                                      placeholder="在此处描述您的需求")
            # 在main()函数的with st.form块中修改：
            col1, col2 = st.columns([3, 1])
            with col1:
                submitted = st.form_submit_button("一键生成测试用例")
            with col2:
                enable_knowledge = st.checkbox("是否参考知识库内容", value=True,
                                               help="启用后将会结合知识库内容生成更精准的用例")

        # 结果生成
        if submitted and user_input:
            with st.spinner("大模型正在努力思考中，请耐心等待哦..."):
                test_cases, sources = generate_test_cases(user_input, model_name=selected_model, enable_knowledge=enable_knowledge)
                if test_cases:
                    st.success(f"✅ 成功生成 {len(test_cases)} 条测试用例！")
                    display_results(test_cases, sources)
                else:
                    st.warning("⚠️ 未生成有效测试用例，请重试")

if __name__ == "__main__":
    main()









